Skip to content

[SPARK-57499][SQL] Fix column pruning and invalid plans in variant extraction pushdown on DSv2 scans - #56556

Closed
qlong wants to merge 2 commits into
apache:masterfrom
qlong:SPARK-57499-variant-pushdown-column-pruning
Closed

[SPARK-57499][SQL] Fix column pruning and invalid plans in variant extraction pushdown on DSv2 scans#56556
qlong wants to merge 2 commits into
apache:masterfrom
qlong:SPARK-57499-variant-pushdown-column-pruning

Conversation

@qlong

@qlong qlong commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

Three fixes in pushVariantExtractions (called byV2ScaqqnRelationPushDown.pushDownVariants):

  1. Guard against double-visit: Add a pushedVariants.isEmpty
    sentinel check so the inner ScanBuilderHolder leaf visit (caused by
    transformDown recursing into the child after returning the plan
    unchanged) returns immediately. This ensures
    builder.pushVariantExtractions is called exactly once per holder.
  2. Eager column pruning: While projectList and filters are in
    scope, call builder.pruneColumns(requiredSchema) for builders
    implementing SupportsPushDownRequiredColumns and trim
    sHolder.output to the required columns. By the time
    buildScanWithPushedVariants calls build(), the builder already
    has the correct pruned schema. This is similiar to how
    buildScanWithPushedAggregate works.
  3. Keep whole-variant reads raw: pushdown is for extractions, not
    whole-variant reads. A bare variant reference -- SELECT v, or a
    column lifted to feed a variant_get above a Join/Sort/Aggregate
    barrier the local rewrite cannot see -- is recorded as fullVariant
    (path $), meaning "the entire value." Shredding that to a lone
    full-variant slot saves no I/O and is mishandled: the Parquet reader
    collapses it to a boolean placeholder, and above a barrier the
    re-exposed GetStructField AS v#orig alias is dropped by
    RemoveRedundantAliases, giving wrong results or an invalid plan. So
    when fullVariant is a variant's only requested field, leave the
    column raw; when it coexists with real extractions (e.g. SELECT v, variant_get(v, '$.a')) the >=2-slot struct is not collapsed and
    keeps its pushdown. This subsumes the join-key case and shreds
    automatically once barrier-aware pushdown makes the variant_get
    visible as a typed path.

Jira: https://issues.apache.org/jira/browse/SPARK-57499

Why are the changes needed?

Three bugs on the accepted variant pushdown path:

Issue 1 (Performance) — column pruning is skipped. buildScanWithPushedVariants calls
builder.build() and replaces the ScanBuilderHolder with a DataSourceV2ScanRelation.
The subsequent pruneColumns rule matches only ScanBuilderHolder nodes, so it is a
no-op and builder.pruneColumns() is never called. The scan reads the full table schema
including unreferenced columns. For unreferenced VARIANT columns this is especially
costly — each is fully reconstructed from its shredded Parquet tree on every row.

Issue 2 (Correctness) — invalid plan / crash on tables with >=2 VARIANT columns
pushDownVariants uses transformDown, which recurses into the child ScanBuilderHolder after returning the plan unchanged. The bare ScanBuilderHolder matches PhysicalOperation a second time, collecting an unreferenced sibling VARIANT column as a full-variant request and pushing it to the builder again. ParquetScanBuilder overwrites its state on every call, so the second push clobbers the correct extraction from the first. The rewritten scan then emits a fresh ExprId for the variant while the projection still references the original, and binding fails.

This affects any extraction shape — projection, ORDER BY, aggregate, join — on a table with two or more VARIANT columns. Single-VARIANT tables are unaffected.

Reproduce on stock spark-4.1.x (path-based views force DSv2):

SET spark.sql.sources.useV1SourceList = "";

CREATE TABLE t (a INT, v1 VARIANT, v2 VARIANT) USING PARQUET LOCATION '/tmp/vt';
INSERT INTO t VALUES
  (1, parse_json('{"x":1,"price":3,"name":"x"}'), parse_json('{"y":2}')),
  (2, parse_json('{"x":9,"price":1,"name":"z"}'), parse_json('{"y":8}'));
CREATE OR REPLACE TEMPORARY VIEW tv    USING parquet OPTIONS (path '/tmp/vt');
CREATE OR REPLACE TEMPORARY VIEW codes USING parquet OPTIONS (path '/tmp/vt');

-- All four crash with: [INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#NN in [...]
SELECT variant_get(v1, '$.x', 'int') FROM tv;        
SELECT variant_get(v1,'$.name','string') AS nm
  FROM tv ORDER BY variant_get(v1,'$.price','int');                   
SELECT max(variant_get(v1, '$.price', 'int')) FROM tv;                     
SELECT l.a FROM tv l JOIN codes r
  ON variant_get(l.v1,'$.x','int') = variant_get(r.v1,'$.x','int');  

[INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#21 in [a#33,v1#34,v2#35]. SQLSTATE: XX000
  at org.apache.spark.sql.catalyst.expressions.BindReferences$.attributeNotFoundException(BoundAttribute.scala:109)
  ...

Issue 3 (Correctness) -- wrong results / crash when a whole-variant read is shredded. A bare
variant reference (a plain SELECT v, or a column lifted to feed a variant_get
above a Join/Sort/Aggregate barrier the local rewrite cannot see) is recorded as a
full-variant request (path "$"). Shredding it to a lone full-variant slot is both
useless (the whole value is read regardless) and mishandled:

  • The Parquet reader collapses a lone VariantType slot to a boolean placeholder,
    so ORDER BY variant_get(v, '$.price') sorts on the placeholder and silently
    returns the wrong order, and max(variant_get(v, '$.price')) fails to codegen.
  • A join key is re-exposed above the join as GetStructField(v_new, i) AS v#orig;
    RemoveRedundantAliases collapses the alias and the condition references a
    dropped ExprId, failing plan validation.

This issue is masked on stock 4.1 by the Issue 2 binding crash, which fails first; it surfaces only after that issue is fixed, which is why it has no repro for stock 4.1. See new unit tests for join, order by, aggregrate.

Does this PR introduce any user-facing change?

No

How was this patch tested?

  • Added new unit tests
  • manual testing with spark-sql

Was this patch authored or co-authored using generative AI tooling?

Co-authored with Claude code (Sonnet 4.6)

@qlong qlong changed the title [SPARK-57499][SQL] Variant extraction pushdown bypasses column pruning on DSvs scans [SPARK-57499][SQL] Variant extraction pushdown bypasses column pruning on DSvs2 scans Jun 17, 2026
@qlong
qlong force-pushed the SPARK-57499-variant-pushdown-column-pruning branch from 5abbc88 to 34f3787 Compare June 17, 2026 16:16
@cloud-fan

Copy link
Copy Markdown
Contributor

can you fix merge conflicts?

@qlong
qlong force-pushed the SPARK-57499-variant-pushdown-column-pruning branch from 34f3787 to 0a76001 Compare June 18, 2026 01:37
@qlong

qlong commented Jun 18, 2026

Copy link
Copy Markdown
Contributor Author

@cloud-fan rebased. I will keep an eye on the CI build

…g on DSv2 scans

Three fixes in `pushVariantExtractions` (called by
`V2ScaqqnRelationPushDown.pushDownVariants`):

1. **Guard against double-visit**: Add a `pushedVariants.isEmpty`
   sentinel check so the inner `ScanBuilderHolder` leaf visit (caused by
   `transformDown` recursing into the child after returning the plan
   unchanged) returns immediately. This ensures
   `builder.pushVariantExtractions` is called exactly once per holder.

2. **Eager column pruning**: While `projectList` and `filters` are in
   scope, call `builder.pruneColumns(requiredSchema)` for builders
   implementing `SupportsPushDownRequiredColumns` and trim
   `sHolder.output` to the required columns. By the time
   `buildScanWithPushedVariants` calls `build()`, the builder already
   has the correct pruned schema. This is similiar to how
   buildScanWithPushedAggregate works.

3. **Keep whole-variant reads raw**: pushdown is for extractions, not
   whole-variant reads. A bare variant reference -- `SELECT v`, or a
   column lifted to feed a `variant_get` above a Join/Sort/Aggregate
   barrier the local rewrite cannot see -- is recorded as `fullVariant`
   (path `$`), meaning "the entire value." Shredding that to a lone
   full-variant slot saves no I/O and is mishandled: the Parquet reader
   collapses it to a boolean placeholder, and above a barrier the
   re-exposed `GetStructField AS v#orig` alias is dropped by
   `RemoveRedundantAliases`, giving wrong results or an invalid plan. So
   when fullVariant is a variant's only requested field, leave the
   column raw; when it coexists with real extractions (e.g. `SELECT v,
   variant_get(v, '$.a')`) the >=2-slot struct is not collapsed and
   keeps its pushdown. This subsumes the join-key case and shreds
   automatically once barrier-aware pushdown makes the `variant_get`
   visible as a typed path.

Jira: https://issues.apache.org/jira/browse/SPARK-57499
@qlong
qlong force-pushed the SPARK-57499-variant-pushdown-column-pruning branch from 0a76001 to 81c15e5 Compare June 18, 2026 05:40
@qlong qlong changed the title [SPARK-57499][SQL] Variant extraction pushdown bypasses column pruning on DSvs2 scans [SPARK-57499][SQL] Fix column pruning and invalid plans in variant extraction pushdown on DSv2 scans Jun 18, 2026

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 blocking, 2 non-blocking, 1 nit.
Solid, well-tested fix; rule structure is correct and consistent with the aggregate path. One correctness gap: the whole-variant guard is narrower than the Parquet reader's placeholder-collapse condition.

Correctness (2)

  • V2ScanRelationPushDown.scala:476: whole-variant strip matches only the exact fullVariant sentinel, so 2-arg variant_get (VariantType target) escapes it and the reader collapses it to a boolean placeholder → wrong results — see inline
  • V2ScanRelationPushDown.scala:373: one early return (!supportsDeferCastError) omits the sentinel this comment promises — see inline

Suggestions (1)

  • PushVariantIntoScanSuite.scala: no 2-arg variant_get (VariantType target) coverage — a case asserting the column is kept raw (e.g. variant_get(v, '$.a'), and variant_get(v, '$') under a non-UTC session) would catch and guard the blocking finding above.

Nits: 1 minor item (see inline comments).

Verification

Traced the variant-extraction rewrite for equivalence. Equivalent for: bare whole reads (collected as the UTC fullVariant → stripped → kept raw, timezone-robust), non-variant-target extractions, IsNull/IsNotNull-only columns (placeholder injection), and NULL / nondeterminism / ordering (variant_get is deterministic, evaluated once over the pushed slot). The one non-equivalent cell: a sole VariantType-target extraction that isn't the exact fullVariantvariant_get(v,'$.a') (2-arg) or variant_get(v,'$') non-UTC — which the reader collapses to Boolean while catalyst keeps VariantType; the top-level-count assert doesn't catch the inner-type mismatch. Output-attribute metadata (nullability/exprId) is consistent — the rewrite synthesizes a Project over the scan's own readSchema and preserves the original column nullability; no outer-join-style widening.

// unaffected.
variants.mapping.values.foreach { pathToFields =>
pathToFields.filterInPlace { case (_, fields) =>
!(fields.size == 1 && fields.contains(RequestedVariantField.fullVariant))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pushVariantExtractions strips a column's mapping only when its sole requested field equals the exact fullVariant sentinel (VariantMetadata("$", failOnError=true, "UTC")). But a 2-arg variant_get builds VariantGet(v, path, VariantType, …) carrying the session timezone and the real path, so variant_get(v, '$.a') (any session) and variant_get(v, '$') (non-UTC session) don't match and escape the strip. They then shred to a lone VariantType slot, which ParquetScan.rewriteVariantPushdownSchema collapses to a BooleanType placeholder — while catalyst's VariantInRelation.rewriteType keeps it VariantType. So holder.output declares struct{VariantType} but the scan produces struct{Boolean}; the new top-level-count assert (1==1) doesn't catch the inner-type mismatch, and GetStructField(0) reads the placeholder → wrong results / invalid plan — Issue 3 via the function form. Aligning the guard with the reader's collapse predicate fixes all of these and subsumes the bare-read case:

Suggested change
!(fields.size == 1 && fields.contains(RequestedVariantField.fullVariant))
!(fields.size == 1 && fields.head._1.targetType.isInstanceOf[VariantType])

Worth adding a 2-arg variant_get test — none exists today, so CI is blind to this.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed, targetType.isInstanceOf[VariantType]) is more direct, added tests for 2-arg.

// Without a guard, this would add fullVariant for *every* variant column.
//
// The `pushedVariants.isEmpty` guard prevents the inner visit from re-running once
// the outer visit has completed (successfully or not). An empty-mapping sentinel

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment says the sentinel is written by the outer visit on completion "successfully or not", but the early return at the hasCompanionExtraction && !builder.supportsDeferCastError() check below (~L554) returns originalPlan without setting sHolder.pushedVariants. That leaves the guard's isEmpty true, so the inner leaf re-visit (Issue 2) would fire. It's masked today — ParquetScanBuilder.supportsDeferCastError() returns true, so the branch is unreachable for the only OSS implementer, and the new full-variant strip would also neutralize a re-visit — but it's a latent reintroduction of the bug for any future builder returning false. Suggest setting sHolder.pushedVariants = Some(new VariantInRelation()) before that return, consistent with the other early returns.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for review. Changed to set sHolder.pushedVariants = Some(new VariantInRelation()) before retturn.

spark.read.parquet(path).createOrReplaceTempView("T_V2")
// variant_get(v, '$.price') is inside an aggregate function, above the aggregate barrier,
// with no local filter/projection on v. v is lifted in only as a bare reference -> a
// whole-variant request. Before the fix this shredded to a lone full-variant slot, which

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: this comment (and a couple of others, e.g. the "SPARK-57499 regression" note) is framed as PR history — "Before the fix this shredded…". Once merged that history is invisible to readers. Consider rephrasing in terms of the invariant the test guards now, e.g. "A whole-variant read is kept raw; shredding it to a lone full-variant slot would collapse to a boolean placeholder and corrupt the result."

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

- Align the whole-variant strip with the reader's collapse predicate:
  match a sole field by VariantType target type instead of the exact
  `fullVariant` value, so 2-arg `variant_get` and non-UTC sessions are
  also kept raw (they shred to the same boolean-placeholder slot).
- Set the pushedVariants sentinel on the
  `hasCompanionExtraction && !supportsDeferCastError()` early return,
  consistent with the other early returns, so the leaf is not re-visited.
- Add tests for 2-arg `variant_get` and `variant_get('$')` under a
  non-UTC session, asserting the column is kept raw.
@qlong
qlong force-pushed the SPARK-57499-variant-pushdown-column-pruning branch from 5bb16e5 to 39899b2 Compare June 22, 2026 03:55
@cloud-fan

Copy link
Copy Markdown
Contributor

thanks, merging to master/4.x

@cloud-fan cloud-fan closed this in bd65ed7 Jun 22, 2026
cloud-fan pushed a commit that referenced this pull request Jun 22, 2026
…traction pushdown on DSv2 scans

### What changes were proposed in this pull request?
Three fixes in `pushVariantExtractions` (called by`V2ScaqqnRelationPushDown.pushDownVariants`):

1. **Guard against double-visit**: Add a `pushedVariants.isEmpty`
       sentinel check so the inner `ScanBuilderHolder` leaf visit (caused by
       `transformDown` recursing into the child after returning the plan
       unchanged) returns immediately. This ensures
       `builder.pushVariantExtractions` is called exactly once per holder.
 2. **Eager column pruning**: While `projectList` and `filters` are in
       scope, call `builder.pruneColumns(requiredSchema)` for builders
       implementing `SupportsPushDownRequiredColumns` and trim
       `sHolder.output` to the required columns. By the time
       `buildScanWithPushedVariants` calls `build()`, the builder already
       has the correct pruned schema. This is similiar to how
       buildScanWithPushedAggregate works.
  3. **Keep whole-variant reads raw**: pushdown is for extractions, not
       whole-variant reads. A bare variant reference -- `SELECT v`, or a
       column lifted to feed a `variant_get` above a Join/Sort/Aggregate
       barrier the local rewrite cannot see -- is recorded as `fullVariant`
       (path `$`), meaning "the entire value." Shredding that to a lone
       full-variant slot saves no I/O and is mishandled: the Parquet reader
       collapses it to a boolean placeholder, and above a barrier the
       re-exposed `GetStructField AS v#orig` alias is dropped by
       `RemoveRedundantAliases`, giving wrong results or an invalid plan. So
       when fullVariant is a variant's only requested field, leave the
       column raw; when it coexists with real extractions (e.g. `SELECT v,
       variant_get(v, '$.a')`) the >=2-slot struct is not collapsed and
       keeps its pushdown. This subsumes the join-key case and shreds
       automatically once barrier-aware pushdown makes the `variant_get`
       visible as a typed path.

Jira: https://issues.apache.org/jira/browse/SPARK-57499

### Why are the changes needed?

Three bugs on the accepted variant pushdown path:

**Issue 1 (Performance) — column pruning is skipped.** `buildScanWithPushedVariants` calls
`builder.build()` and replaces the `ScanBuilderHolder` with a `DataSourceV2ScanRelation`.
The subsequent `pruneColumns` rule matches only `ScanBuilderHolder` nodes, so it is a
no-op and `builder.pruneColumns()` is never called. The scan reads the full table schema
including unreferenced columns. For unreferenced `VARIANT` columns this is especially
costly — each is fully reconstructed from its shredded Parquet tree on every row.

**Issue 2 (Correctness) — invalid plan / crash on tables with >=2 VARIANT columns**
pushDownVariants uses transformDown, which recurses into the child ScanBuilderHolder after returning the plan unchanged. The bare ScanBuilderHolder matches PhysicalOperation a second time, collecting an unreferenced sibling VARIANT column as a full-variant request and pushing it to the builder again. ParquetScanBuilder overwrites its state on every call, so the second push clobbers the correct extraction from the first. The rewritten scan then emits a fresh ExprId for the variant while the projection still references the original, and binding fails.

This affects any extraction shape — projection, ORDER BY, aggregate, join — on a table with two or more VARIANT columns. Single-VARIANT tables are unaffected.

Reproduce on stock spark-4.1.x (path-based views force DSv2):

```
SET spark.sql.sources.useV1SourceList = "";

CREATE TABLE t (a INT, v1 VARIANT, v2 VARIANT) USING PARQUET LOCATION '/tmp/vt';
INSERT INTO t VALUES
  (1, parse_json('{"x":1,"price":3,"name":"x"}'), parse_json('{"y":2}')),
  (2, parse_json('{"x":9,"price":1,"name":"z"}'), parse_json('{"y":8}'));
CREATE OR REPLACE TEMPORARY VIEW tv    USING parquet OPTIONS (path '/tmp/vt');
CREATE OR REPLACE TEMPORARY VIEW codes USING parquet OPTIONS (path '/tmp/vt');

-- All four crash with: [INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#NN in [...]
SELECT variant_get(v1, '$.x', 'int') FROM tv;
SELECT variant_get(v1,'$.name','string') AS nm
  FROM tv ORDER BY variant_get(v1,'$.price','int');
SELECT max(variant_get(v1, '$.price', 'int')) FROM tv;
SELECT l.a FROM tv l JOIN codes r
  ON variant_get(l.v1,'$.x','int') = variant_get(r.v1,'$.x','int');

[INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find v1#21 in [a#33,v1#34,v2#35]. SQLSTATE: XX000
  at org.apache.spark.sql.catalyst.expressions.BindReferences$.attributeNotFoundException(BoundAttribute.scala:109)
  ...
```

**Issue 3 (Correctness) -- wrong results / crash when a whole-variant read is shredded**. A bare
variant reference (a plain `SELECT v`, or a column lifted to feed a variant_get
above a Join/Sort/Aggregate barrier the local rewrite cannot see) is recorded as a
full-variant request (path "$"). Shredding it to a lone full-variant slot is both
useless (the whole value is read regardless) and mishandled:
  - The Parquet reader collapses a lone VariantType slot to a boolean placeholder,
    so ORDER BY variant_get(v, '$.price') sorts on the placeholder and silently
    returns the wrong order, and max(variant_get(v, '$.price')) fails to codegen.
  - A join key is re-exposed above the join as GetStructField(v_new, i) AS v#orig;
    RemoveRedundantAliases collapses the alias and the condition references a
    dropped ExprId, failing plan validation.

This issue is masked on stock 4.1 by the Issue 2 binding crash, which fails first; it surfaces only after that issue is fixed, which is why it has no repro for stock 4.1. See new unit tests for join, order by, aggregrate.

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?

- Added new unit tests
- manual testing with spark-sql

### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude code (Sonnet 4.6)

Closes #56556 from qlong/SPARK-57499-variant-pushdown-column-pruning.

Authored-by: Qiegang Long <qlong@users.noreply.github.com>
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
(cherry picked from commit bd65ed7)
Signed-off-by: Wenchen Fan <wenchen@databricks.com>
@qlong

qlong commented Jun 22, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for review and merge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants